import math
import numpy as np

import pat.autopilot as ap
import pat.vehicles.rotorcraft.dynamic_model_hk as dm
import pat.utils as pu

class Autopilot(ap.Autopilot):
    def __init__(self, _dm, _args=None, _time=None):
        ap.Autopilot.__init__(self)
        self.add_mode(AttLaw(_dm))
        self.add_mode(PositionLaw(_dm))
        self.set_cur_mode(1)
        
    def get_U_str(self):
        return "{:-.3f} {:-.3f} {:-.3f} {:-.3f}".format(self.U[dm.iv_pitch], self.U[dm.iv_roll], self.U[dm.iv_col], self.U[dm.iv_yaw])
    
class Step_x(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, dm.sv_x, "Step x") # x
    
class Step_z(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, dm.sv_z, "Step z") # x

class Step_phi(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, dm.sv_phi, "Step phi", ampl=pu.rad_of_deg(5)) # x


class AttLaw(ap.ControlLaw):
    def __init__(self, _dm):
        ap.ControlLaw.__init__(self, "Att")
        self.K0_pil = np.array([ 0.0 ,   0.0 ,   0.1])
        self.Kp_pil = np.array([-0.3  ,  0.25 , -0.3])
        self.Kd_pil = np.array([-0.05 ,  0.05 , -0.05])
        self.Ki_pil = np.array([ 0.0,    0.0,   -0.0002])
        self.sum_err_pil = np.zeros(3)
        self.U = np.zeros(dm.iv_size)
        Xe, Ue = _dm.trim()
        self.reference= [Step_phi(Xe, Ue)]

    def run(self, time, X, ref):
        euler = X[dm.sv_phi:dm.sv_psi+1]; rvel  = X[dm.sv_p:dm.sv_r+1]
        euler_sp = [pu.rad_of_deg(3.), 0., 0.]
        if ref:
            Xref, Uref = ref.get(time)
            euler_sp += Xref[dm.sv_phi:dm.sv_psi+1]
            
        euler_err = euler - euler_sp
        euler_err[2] = norm_psi(euler_err[2])
        self.sum_err_pil += euler_err
        roll, pitch, yaw = self.K0_pil + self.Kp_pil*euler_err + self.Kd_pil*rvel + self.Ki_pil*self.sum_err_pil
        coll = 0.16024332
        U = (roll, pitch, yaw, coll)
        self.U = np.array([pitch, roll, coll, yaw])
#        self.U = np.array([-0.00,  0.00,  0.16024332,  0.10031169])
        return self.U, ref

    def properties(self):
        return "pil Kp:{:s}\npil Kd:{:s}\npil Ki:{:s}".format(self.Kp_pil, self.Kd_pil, self.Ki_pil)


        

class PositionLaw(ap.ControlLaw):
    def __init__(self, _dm):
        ap.ControlLaw.__init__(self, "Position")
        self.K0_guid = np.array([0.0 ,     0.0 ,    0.16])
        self.Kp_guid = np.array([0.015,   -0.015,   0.01])
        self.Kd_guid = np.array([0.02,    -0.02,    0.01])       
        self.Ki_guid = np.array([0.00005, -0.00005, 0.00001])
        self.sum_err_guid = np.zeros(3)

        self.K0_pil = np.array([ 0.0 ,   0.0 ,   0.1])
        self.Kp_pil = np.array([-0.25 ,  0.25 , -0.3])
        self.Kd_pil = np.array([-0.05 ,  0.05 , -0.05])
        self.Ki_pil = np.array([ 0.0,    0.0,   -0.0002])
        self.sum_err_pil = np.zeros(3)

        self.U = np.zeros(dm.iv_size)

        self.pos_sp   = np.array([0., 0., -4.])
        self.vel_sp   = np.array([0., 0.,  0.])
        self.euler_sp = np.array([pu.rad_of_deg(3.), 0., 0.])

        self.Kp_guid[0] =  0.02
        self.Kd_guid[0] =  0.02


        self.Kp_guid[1] = -0.04
        self.Kd_guid[1] = -0.3

        self.Kp_pil[0] = -0.25
        self.Kd_pil[0] = -0.05

        Xe, Ue = _dm.trim()
        self.reference= [Step_x(Xe, Ue), Step_z(Xe, Ue)]

        
    def run(self, time, X, ref):
        pos   = X[dm.sv_x:dm.sv_z+1]; vel   = X[dm.sv_vx:dm.sv_vz+1]
        euler = X[dm.sv_phi:dm.sv_psi+1]; rvel  = X[dm.sv_p:dm.sv_r+1]

        if ref:
            Xref, Uref = ref.get(time)
            pos_ref = Xref[dm.sv_x:dm.sv_z+1]
            vel_ref = Xref[dm.sv_x:dm.sv_z+1]
        else:
            pos_ref = self.pos_sp
            vel_ref = self.vel_sp
        cpsi = math.cos(euler[2]); spsi =  math.sin(euler[2])
        rmat = np.array([[cpsi, spsi, 0.], [-spsi, cpsi, 0.], [0., 0., 1.]])
        pos_err_ned = pos - pos_ref 
        pos_err_wa = np.dot(rmat, pos_err_ned)
        self.sum_err_guid += pos_err_wa

        vel_err_ned = vel - vel_ref
        vel_err_wa = np.dot(rmat, vel_err_ned)
        
        pitch_sp, roll_sp, coll = self.K0_guid + self.Kp_guid*pos_err_wa + self.Kd_guid*vel_err_wa + self.Ki_guid*self.sum_err_guid

        euler_sp = self.euler_sp + np.array([roll_sp, pitch_sp, 0.])
        euler_err = euler - euler_sp
        euler_err[2] = norm_psi(euler_err[2])
        self.sum_err_pil += euler_err
        roll, pitch, yaw = self.K0_pil + self.Kp_pil*euler_err + self.Kd_pil*rvel + self.Ki_pil*self.sum_err_pil
        U = (roll, pitch, yaw, coll)
        self.U = np.array([pitch, roll, coll, yaw])
#        self.U = np.array([-0.00,  0.00,  0.16024332,  0.10031169])
#        print self.U
        return self.U, ref

    def properties(self):
        
        s1 = "pil Kp:{:s}\npil Kd:{:s}\npil Ki:{:s}\n".format(self.Kp_pil, self.Kd_pil, self.Ki_pil)
        s2 = "guid Kp:{:s}\nguid Kd:{:s}\nguid Ki:{:s}".format(self.Kp_guid, self.Kd_guid, self.Ki_guid)
        return s1+s2



def norm_psi(psi):
    psi_n = math.fmod(psi, 2*math.pi)
    if psi_n > math.pi: psi_n = psi_n - 2*math.pi
    return psi_n
